use crate::{
    disk::{HD_LIST, ahci::AHCI},
    ref_to_mut,
};
use alloc::{
    string::{String, ToString},
    sync::Arc,
    vec::Vec,
};
use spin::RwLock;
use syscall::scheme::{KernelScheme, SchemeRef};

use super::{
    ROOT_INODE,
    cache::{BlockDeviceInterface, Cache512B, CacheManager},
    get_scheme_by_path,
    gpt_parser::parse_gpt_disk,
    mount_to,
};

struct BlockDevice {
    id: usize,
}

impl BlockDevice {
    pub fn new(id: usize) -> Self {
        Self { id }
    }
}

impl BlockDeviceInterface for BlockDevice {
    fn read(&self, block_id: usize, buf: &mut [u8]) {
        HD_LIST.lock()[self.id].read_block(block_id, buf);
    }

    fn write(&self, block_id: usize, buf: &[u8]) {
        HD_LIST.lock()[self.id].write_block(block_id, buf);
    }
}

pub struct BlockScheme {
    hd: usize,
    cache_manager: CacheManager<Cache512B, BlockDevice>,
    path: String,
}

impl BlockScheme {
    pub fn new(hd: usize) -> Self {
        Self {
            hd,
            cache_manager: CacheManager::new(BlockDevice::new(hd)),
            path: String::new(),
        }
    }
}

impl KernelScheme for BlockScheme {
    fn when_mounted(&mut self, path: alloc::string::String, _father: Option<SchemeRef>) {
        self.path.clear();
        self.path.push_str(path.as_str());
    }

    fn when_umounted(&mut self) {}

    fn get_path(&self) -> alloc::string::String {
        self.path.clone()
    }

    fn size(&self) -> usize {
        HD_LIST.lock()[self.hd].get_size()
    }

    fn read_at(&self, offset: usize, buf: &mut [u8]) -> usize {
        let start = offset;
        let end = start + buf.len();

        let start_sector_read_start = start % 512;

        let start_sector_id = start / 512;
        let end_sector_id = (end - 1) / 512;

        let buffer_size = (end_sector_id - start_sector_id + 1) * 512;
        let mut tmp = Vec::new();
        for _ in 0..buffer_size {
            tmp.push(0);
        }
        let tmp = tmp.leak();

        ref_to_mut(self)
            .cache_manager
            .read_from_cache(start_sector_id, tmp);

        for i in 0..(end - start) {
            buf[i] = tmp[i + start_sector_read_start];
        }
        buf.len()
    }

    fn write_at(&self, offset: usize, buf: &[u8]) -> usize {
        let start = offset;
        let end = start + buf.len();

        let start_sector_read_start = start % 512;

        let start_sector_id = start / 512;
        let end_sector_id = (end - 1) / 512;

        let buffer_size = (end_sector_id - start_sector_id + 1) * 512;
        let mut tmp = Vec::new();
        for _ in 0..buffer_size {
            tmp.push(0);
        }
        let tmp = tmp.leak();

        ref_to_mut(self)
            .cache_manager
            .read_from_cache(start_sector_id, tmp);

        for i in 0..(end - start) {
            tmp[i + start_sector_read_start] = buf[i];
        }

        ref_to_mut(self)
            .cache_manager
            .write_to_cache(start_sector_id, tmp);

        ref_to_mut(self).cache_manager.flush_cache();

        buf.len()
    }
}

struct AHCIDisk {
    num: usize,
}

pub trait BlockDevicea: Send + Sync + 'static {
    fn read_block(&self, start_sec: usize, buf: &mut [u8]) -> Option<()>;
    fn write_block(&self, start_sec: usize, buf: &[u8]) -> Option<()>;

    fn get_size(&self) -> usize;
}

impl BlockDevicea for AHCIDisk {
    fn read_block(&self, start_sec: usize, buf: &mut [u8]) -> Option<()> {
        AHCI.lock()
            .get_disk(self.num)
            .expect("Cannot get disk")
            .read_block(start_sec as u64, buf);
        Some(())
    }

    fn write_block(&self, start_sec: usize, buf: &[u8]) -> Option<()> {
        AHCI.lock()
            .get_disk(self.num)
            .expect("Cannot get disk")
            .write_block(start_sec as u64, buf);
        Some(())
    }

    fn get_size(&self) -> usize {
        AHCI.lock()
            .get_disk(self.num)
            .unwrap()
            .get_identity()
            .lba48_sectors as usize
            * 512
    }
}

pub fn init() {
    let ahci_disk_num = AHCI.lock().len();

    for num in 0..ahci_disk_num {
        let disk = Arc::new(AHCIDisk { num });
        HD_LIST.lock().push(disk.clone());

        let blk_scheme = Arc::new(RwLock::new(BlockScheme::new(num)));
        let id_to_alpha = [
            "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q",
            "r", "s", "t", "u", "v", "w", "x", "y", "z",
        ];

        mount_to(
            blk_scheme.clone(),
            ROOT_INODE.lock().clone(),
            alloc::format!("hd{}", id_to_alpha[num]),
        );

        parse_gpt_disk(
            num,
            blk_scheme.clone(),
            get_scheme_by_path("/dev".to_string()).unwrap(),
        )
        .expect("Cannot parse disk")
    }
}
